package org.aeon.fasttasks.core.config;

import org.aeon.fasttasks.core.properties.TaskProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.util.concurrent.RejectedExecutionException;


/**
 * @Date 2024/9/10 16:43
 * @Author aeon
 * @Description 默认的任务执行线程池配置
 */
@Configuration
public class ThreadPoolConfig {

    private final static Logger log = LoggerFactory.getLogger(ThreadPoolConfig.class);

    @Bean("batchWorkTasksExecutor")
    @ConditionalOnMissingBean(name = "batchWorkTasksExecutor")
    public ThreadPoolTaskExecutor batchWorkTasksExecutor(TaskProperties taskProperties) {
        TaskProperties.BatchWorkTasksExecutorProperties
                batchWorkTasksExecutorProperties = taskProperties.batchWorkTasksExecutorProperties();

        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();

        taskExecutor.setCorePoolSize(batchWorkTasksExecutorProperties.corePoolSize());
        taskExecutor.setMaxPoolSize(batchWorkTasksExecutorProperties.maxPoolSize());
        taskExecutor.setQueueCapacity(batchWorkTasksExecutorProperties.queueCapacity());
        taskExecutor.setKeepAliveSeconds(batchWorkTasksExecutorProperties.keepAliveSeconds());
        taskExecutor.setThreadNamePrefix(batchWorkTasksExecutorProperties.threadNamePrefix());
        taskExecutor.setAllowCoreThreadTimeOut(batchWorkTasksExecutorProperties.allowCoreThreadTimeOut());

        // 确保任务不会丢失
        taskExecutor.setRejectedExecutionHandler((r, e) -> {
            if (!e.isShutdown()) {
                log.warn("任务队列已满，等待重新加入队列...");
                try {
                    e.getQueue().put(r);
                } catch (InterruptedException ex) {
                    throw new RuntimeException(ex);
                }
            }
        });

        // 线程池初始化
        taskExecutor.initialize();
        return taskExecutor;
    }


    @Bean("bossTasksExecutor")
    @ConditionalOnMissingBean(name = "bossTasksExecutor")
    public ThreadPoolTaskExecutor bossTasksExecutor(TaskProperties taskProperties) {
        TaskProperties.BossTasksExecutorProperties
                bossTasksExecutorProperties = taskProperties.bossTasksExecutorProperties();

        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setCorePoolSize(bossTasksExecutorProperties.corePoolSize());
        taskExecutor.setMaxPoolSize(bossTasksExecutorProperties.maxPoolSize());
        taskExecutor.setQueueCapacity(0);
        taskExecutor.setThreadNamePrefix("boss-thread-");

        taskExecutor.setRejectedExecutionHandler((r, executor) -> {
            log.warn("任务被拒绝，需要调整 bossTasksExecutor 线程池大小...");
            throw new RejectedExecutionException("BossTasksExecutor is full");
        });

        taskExecutor.initialize();
        return taskExecutor;
    }


    @Bean("tasksScheduler")
    @ConditionalOnMissingBean(name = "tasksScheduler")
    public ThreadPoolTaskScheduler taskScheduler() {
        ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
        scheduler.setPoolSize(1);
        scheduler.setThreadNamePrefix("scheduler-");
        scheduler.initialize();
        return scheduler;
    }

}
